from ultralytics import YOLO
from config.settings import DEFAULT_CONF_THRESHOLD, DEFAULT_IOU_THRESHOLD


class YOLONet:
    def __init__(self):
        self.model = None
        self.conf_threshold = DEFAULT_CONF_THRESHOLD
        self.iou_threshold = DEFAULT_IOU_THRESHOLD

    def load_model(self, model_path):
        """加载YOLOv8模型"""
        try:
            self.model = YOLO(model_path)
            return True, "模型加载成功"
        except Exception as e:
            return False, f"模型加载失败: {str(e)}"

    def detect(self, image_path):
        """执行目标检测"""
        if not self.model:
            return None, 0, {"error": "模型未加载"}

        results = self.model.predict(
            source=image_path,
            conf=self.conf_threshold,
            iou=self.iou_threshold,
            verbose=False
        )

        result = results[0]
        detected_image = result.plot()
        sheep_count = len(result.boxes)

        metrics = {
            "confidence_min": float(result.boxes.conf.min()) if sheep_count > 0 else 0,
            "confidence_max": float(result.boxes.conf.max()) if sheep_count > 0 else 0,
            "confidence_avg": float(result.boxes.conf.mean()) if sheep_count > 0 else 0
        }

        return detected_image, sheep_count, metrics